#!/usr/bin/env python3

import numpy as np
from uarm.wrapper import SwiftAPI
from time import sleep
from vision import transform
import math


class Arm_controller:
    def __init__(self, mtx=None, end2eye_R=None, end2eye_t=None):
        """

        :param mtx: 相机内参
        :param end2eye_R: 手眼旋转矩阵3x3
        :param end2eye_t: 手眼平移矩阵1x3
        """
        # device_info = self.control.get_device_info()
        # print(device_info)
        self.mtx = mtx
        self.end2eye_R = end2eye_R
        self.end2eye_t = end2eye_t
        self.tf = transform.Transform(mtx)

    def arm_init(self):
        self.control = SwiftAPI()
        self.control.reset()


    def move(self,position,speed=None,wait=True,offs=[0,0,0],end_angle=None):
        """
        控制手臂移动
        :param position: xyz坐标
        :param speed: 速度，默认为上一个移动的速度
        :param wait: 是否等待移动完成，默认为True
        :param offs: xyz偏移坐标
        :return: wait为ture时返回控制结果
        """
        if end_angle is not None:
            self.control.set_servo_angle(servo_id=3, angle=end_angle, wait=True)
            sleep(0.5)
        ret = self.control.set_position(position[0]+offs[0],position[1]+offs[1],int(position[2]+offs[2]),speed,wait=wait)

        if ret == "OK":
            # print("Control success")
            return True
        else:
            print("[error]Control failure: ",ret)
            return False

    def compute_point_angle(self,point):
        point = np.array(point[:2])
        angle = 90 + np.arcsin(point[1] / np.sqrt(np.sum(point * point.T))) * 180 / np.pi
        return angle

    def compute_offset_angle(self,start,end,offs=0):
        if offs == 0:
            start_offs = self.compute_point_angle(end)
            end_offs = self.compute_point_angle(start)
        elif offs >= -180 and offs <=180:
            err_angle = offs + self.compute_point_angle(end) - self.compute_point_angle(start)
            if err_angle > 0:
                if err_angle <= 180:
                    start_offs = 180
                    end_offs = 180 - err_angle
                else:
                    start_offs = 0
                    end_offs = 360 - err_angle
            elif err_angle < 0 and err_angle >= -180:
                if err_angle >= -180:
                    start_offs = 0
                    end_offs = -err_angle
                else:
                    start_offs = 180
                    end_offs = 180 + err_angle
        else:
            print(f"输入角度不在-180~180之间")
        return start_offs,end_offs


    def gate_move(self,start,end,offset_z=10,offset_angle=0,speed=None,pump=True):
        """
        门移动指令
        :param start: 起始点
        :param end: 结束点
        :param z_offs: 入刀点
        :param speed: 速度
        :return:
        """
        start_offs_angle,end_offs_angle = self.compute_offset_angle(start,end,offs=offset_angle)
        print(start_offs_angle,end_offs_angle)
        if pump is True:
            self.control.set_pump()
        self.move(start, speed=speed, wait=True, offs=[0, 0, offset_z])

        self.control.waiting_ready()
        self.move(start, speed=speed, wait=True,end_angle=start_offs_angle)
        if pump is True:
            self.control.set_pump(on=True, wait=True)
        sleep(0.5)
        self.move(start, speed=speed, wait=True, offs=[0, 0, offset_z])
        self.move(end, speed=speed, wait=True, offs=[0, 0, offset_z])
        self.control.waiting_ready()
        self.move(end, speed=speed, wait=True, end_angle=end_offs_angle)
        sleep(0.5)
        if pump is True:
            self.control.set_pump(on=False, wait=True)
        self.move(end, speed=speed, wait=True, offs=[0, 0, offset_z])
        return True

    def get_base2end(self):
        R = self.tf.eulerAngles2RotationMatrix([0, 0, (self.control.get_polar()[1] - 90) * math.pi / 180])
        t = self.control.get_position()
        return R,t

    def tag_point(self,image,tag_id,tag_size):
        tags = self.tf.get_img2tag(image,tag_size)
        dst = []
        for tag in tags:
            if tag["tag_id"] == tag_id:
                base2end_R,base2end_t = self.get_base2end()
                goal = self.tf.get_base2tag_p(tag, self.end2eye_R, self.end2eye_t,base2end_R,base2end_t)
                dst.append([goal[0],goal[1],int(goal[2]),tag["angle"]])
        return dst

    def img_point(self,u, v, h):
        base2end_R, base2end_t = self.get_base2end()
        depth = base2end_t - self.end2eye_R @ self.end2eye_t
        img2obj = np.linalg.inv(self.mtx) @ np.array([u, v, 1]) * (depth[2] - h)
        dst = base2end_t + base2end_R @ self.end2eye_R @ img2obj + base2end_R @ -self.end2eye_R @ self.end2eye_t
        return dst

    def make_warehouse(self,point,size,step,angle=0):
        objp = np.mgrid[0:size[0], 0:size[1], 0:size[2]].T.reshape(-1, 3)
        objp = objp*np.array(step)+point
        return objp








